package personal.leil.learning.kruskal;

import personal.leil.learning.undigraph.EdgeUndirected;
import personal.leil.learning.undigraph.Undigraph;
import personal.leil.learning.undigraph.Vertex;

import java.util.*;

/**
 * 最小生成树算法(Kruskal)
 * <p>
 * <pre>
 *         问题：
 *         考虑一个带权重的无向连通图，每个边都有一个与之关联的权重（可以理解为边的成本或距离）。
 *         最小生成树是一个包含所有顶点的树，使得树中的边的权重之和最小。
 *     </pre>
 * <pre>
 *         算法的基本思想：
 *         1.将图中的所有边按照权重进行升序排序。
 *         2.从最小权重的边开始，依次加入生成树，确保加入的边不形成环路。
 *         3.重复步骤2，直到生成树包含了图中的所有顶点为止。
 *     </pre>
 * </p>
 *
 * @author Leil
 * @version 1.0
 * @since 2024/1/23
 */
public class KruskalMST {


    public static List<EdgeUndirected> findMinimumSpanningTree(Undigraph graph) {
        List<EdgeUndirected> list = new ArrayList<>();
        PriorityQueue<EdgeUndirected> priorityQueue = new PriorityQueue<>();
        for (EdgeUndirected edge : graph.edges) {
            priorityQueue.offer(edge);
        }

        /*
         * 创建哈希表，用于存储并查集数据结构的父节点
         * 节点 《=》 父
         */
        Map<Vertex, Vertex> parents = new HashMap<>();
        for (Vertex vertex : graph.vertices) {
            parents.put(vertex, vertex);
        }
        while (!priorityQueue.isEmpty()) {
            EdgeUndirected poll = priorityQueue.poll();
            Vertex[] vertices = poll.vertices;
            /*
             * 使用并查集算法查看两个点有没有相同的parent
             * 如果有相同的parent，则再次连接会产生环路
             */
            Vertex p1 = find(vertices[0], parents);
            Vertex p2 = find(vertices[1], parents);
            if (p1 != p2) {
                union(p1, p2, parents);
                list.add(poll);
            }
        }
        return list;
    }

    /**
     * Find parent
     * @param vertex    target vertex
     * @return          parent
     */
    public static Vertex find(Vertex vertex, Map<Vertex, Vertex> parents) {
        Vertex parent = parents.get(vertex);
        if (vertex == parent) {
            return parent;
        }
        return find(parent, parents);
    }

    /**
     * Union for two trees
     * @param v1        parent of tree1
     * @param v2        parent of tree2
     * @param parents   hash map of parent
     */
    public static void union(Vertex v1, Vertex v2, Map<Vertex, Vertex> parents) {
        parents.put(v1, v2);
    }


    /************************* For Test *************************/
    public static void main(String[] args) {
        Undigraph graph = new Undigraph();
        Vertex v1 = new Vertex(1);
        graph.addVertex(v1);
        Vertex v2 = new Vertex(2);
        graph.addVertex(v2);
        Vertex v3 = new Vertex(3);
        graph.addVertex(v3);
        Vertex v4 = new Vertex(4);
        graph.addVertex(v4);

        v1.addNeighbor(v2);
        graph.addEdge(new EdgeUndirected(v1, v2, 5));
        v2.addNeighbor(v4);
        graph.addEdge(new EdgeUndirected(v2, v4, 2));
        v1.addNeighbor(v3);
        graph.addEdge(new EdgeUndirected(v1, v3, 3));
        v3.addNeighbor(v4);
        graph.addEdge(new EdgeUndirected(v3, v4, 100));
        v1.addNeighbor(v3);
        graph.addEdge(new EdgeUndirected(v1, v4, 9));
        v3.addNeighbor(v2);
        graph.addEdge(new EdgeUndirected(v3, v2, 8));

        List<EdgeUndirected> tree = findMinimumSpanningTree(graph);
        for (EdgeUndirected edge : tree) {
            System.out.println(edge);
        }

    }


}
